import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad
from scipy.optimize import minimize_scalar

# 定义常量
p = 1.7  # 螺距，单位：米
b = p / (2 * np.pi)  # 转换常数
r_turn = 4.5  # 调头区域半径，单位：米
v_head = 1.0  # 龙头速度，单位：米/秒


# 定义计算两个极角之间距离的函数
def polar_distance(theta1, theta2, b):
    r1 = b * theta1
    r2 = b * theta2
    x1 = r1 * np.cos(theta1)
    y1 = r1 * np.sin(theta1)
    x2 = r2 * np.cos(theta2)
    y2 = r2 * np.sin(theta2)
    return np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)


# 定义计算路径长度的函数
def path_length(theta, b):
    r = b * theta
    return np.sqrt(r ** 2 + b ** 2)


# 定义优化目标函数
def objective(r_turn, b):
    # 计算盘入路径长度
    length_in = quad(path_length, 0, r_turn / b)[0]
    # 计算调头路径长度
    # 假设调头路径为圆弧，半径为r_turn
    length_turn = r_turn * np.pi
    # 计算盘出路径长度
    length_out = quad(path_length, r_turn / b, 2 * np.pi * r_turn / b)[0]
    # 返回总路径长度
    return length_in + length_turn + length_out


# 优化调头半径以最小化总路径长度
res = minimize_scalar(objective, args=(b,), bounds=(1, 10), method='bounded')
optimal_r_turn = res.x

print(f"Optimal turning radius: {optimal_r_turn} meters")

# 绘制路径
theta_vals = np.linspace(0, 2 * np.pi * r_turn / b, 1000)
r = b * theta_vals
x = r * np.cos(theta_vals)
y = r * np.sin(theta_vals)

plt.figure()
plt.plot(x, y, label='Path')
plt.scatter(0, 0, color='red', label='Start/End Point')
plt.title('Optimal Path for Dragon Dance')
plt.xlabel('X Position (m)')
plt.ylabel('Y Position (m)')
plt.legend()
plt.axis('equal')
plt.grid(True)
plt.show()
